import numpy as np
from scipy import stats
from pykalman import KalmanFilter
from qtorch.strategy import Strategy

class KalmanFilterStrategy(Strategy):
    """卡尔曼滤波策略（动态回归）"""
    def __init__(self, delta=1e-5, window=30):
        self.delta = delta
        self.window = window
        
    def generate_signals(self, prices):
        if isinstance(prices, pd.DataFrame):
            x = prices['close'].values
            y = prices['open'].values
        else:
            x = prices.values
            y = np.roll(prices.values, 1)
            
        # 初始化卡尔曼滤波器
        kf = KalmanFilter(
            initial_state_mean=0,
            initial_state_covariance=1,
            transition_matrices=[1],
            observation_matrices=[1],
            observation_covariance=1,
            transition_covariance=self.delta
        )
        
        # 运行卡尔曼滤波
        means, _ = kf.filter(y)
        
        # 计算残差和标准差
        residuals = x - means.flatten()
        std = np.sqrt(np.mean(residuals**2))
        
        # 生成信号
        signals = np.where(residuals > std, -1,
                         np.where(residuals < -std, 1, 0))
        return signals.astype(int)